from models.standard_blocks import *


class FengyongNet(nn.Module):
    def __init__(self, num_categories):
        super(FengyongNet, self).__init__()
        self.num_categories = num_categories
        self.conv_blocks = nn.Sequential(half_hw_convblock(1, 16),
                                         equal_hw_convblock(16, 32),
                                         equal_hw_convblock(32, 16),
                                         half_hw_convblock(16, 8),
                                         half_hw_convblock(8, 8))
        self.tail_conv = equal_hw_convblock(8, self.num_categories)

    def tail(self, x):
        x = self.tail_conv(x)
        x = nn.AdaptiveAvgPool2d((1, 1))(x)
        x = x.view(-1, self.num_categories)
        return x

    def forward(self, x):
        x = self.conv_blocks(x)
        x = self.tail(x)

        return x
